package com.niodata.dp.monitor.core.task.entity;

import java.util.PriorityQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

public class KeyCache {

  private static ScheduledExecutorService swapExpiredPool
        = new ScheduledThreadPoolExecutor(Runtime.getRuntime().availableProcessors() * 2);
  private ReentrantLock lock = new ReentrantLock();
  private ConcurrentHashMap<String, Node> cache = new ConcurrentHashMap<>(1024);
  private PriorityQueue<Node> expireQueue = new PriorityQueue<>(1024);

  public KeyCache() {

    //使用默认的线程池，每20秒清除一次过期数据
    swapExpiredPool.scheduleWithFixedDelay(
          new SwapExpiredNodeWork(), 20, 20, TimeUnit.SECONDS);
  }

  public Object set(String key, long ttl) {

    Assert.isTrue(StringUtils.hasLength(key), "key can't be empty");
    Assert.isTrue(ttl > 0, "ttl must greater than 0");

    long expireTime = System.currentTimeMillis() + ttl;
    Node newNode = new Node(key, expireTime);
    lock.lock();
    try {
      Node old = cache.put(key, newNode);
      expireQueue.add(newNode);
      //如果该key存在数据，从过期时间队列删除
      if (old != null) {
        expireQueue.remove(old);
        return old.key;
      }
      return null;
    } finally {
      lock.unlock();
    }

  }

  public Object get(String key) {
    Node n = cache.get(key);
    return n == null ? null : n.key;
  }

  /**
   * 删除KEY，并返回该key的值.
   */
  public Object remove(String key) {
    lock.lock();
    try {
      Node n = cache.remove(key);
      if (n == null) {
        return null;
      } else {
        expireQueue.remove(n);
        return n.key;
      }
    } finally {
      lock.unlock();
    }
  }

  private static class Node implements Comparable<Node> {

    private String key;
    private long expireTime;

    public Node(String key, long expireTime) {
      this.key = key;
      this.expireTime = expireTime;
    }


    /**
     * @see SwapExpiredNodeWork
     */
    @Override
    public int compareTo(Node o) {
      long r = this.expireTime - o.expireTime;
      if (r > 0) {
        return 1;
      }
      if (r < 0) {
        return -1;
      }
      return 0;
    }

  }

  /**
   * 删除已经过期的key .
   */
  private class SwapExpiredNodeWork implements Runnable {

    @Override
    public void run() {

      long now = System.currentTimeMillis();
      while (true) {
        lock.lock();
        try {
          Node node = expireQueue.peek();
          if (node == null || node.expireTime > now) {
            return;
          }
          cache.remove(node.key);
          expireQueue.poll();
        } finally {
          lock.unlock();
        }
      }
    }
  }
}
